import json
import os
import argparse

from pprint import pprint


def attribute_statistics():
    input_path = opt.input_path
    count = dict()

    for dirpath, _, filenames in os.walk(input_path):
        for file in filenames:
            full_file_path = os.path.join(dirpath, file)
            if os.path.isfile(full_file_path):
                with open(full_file_path, "r", encoding="utf-8") as label_file:
                    label = json.load(label_file)
                    attributes = label["attributes"]
                    if "attributes" not in count:
                        count["attributes"] = dict()

                    weather = attributes["weather"]
                    if "weather" not in count["attributes"]:
                        count["attributes"]["weather"] = dict()
                    if weather not in count["attributes"]["weather"]:
                        count["attributes"]["weather"][weather] = 1
                    else:
                        count["attributes"]["weather"][weather] = count["attributes"]["weather"][weather] + 1

                    scene = attributes["scene"]
                    if "scene" not in count["attributes"]:
                        count["attributes"]["scene"] = dict()
                    if scene not in count["attributes"]["scene"]:
                        count["attributes"]["scene"][scene] = 1
                    else:
                        count["attributes"]["scene"][scene] = count["attributes"]["scene"][scene] + 1

                    timeofday = attributes["timeofday"]
                    if "timeofday" not in count["attributes"]:
                        count["attributes"]["timeofday"] = dict()
                    if timeofday not in count["attributes"]["timeofday"]:
                        count["attributes"]["timeofday"][timeofday] = 1
                    else:
                        count["attributes"]["timeofday"][timeofday] = count["attributes"]["timeofday"][timeofday] + 1
    return count

def dict_save2json(data, path):
    output_dir = os.path.dirname(path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"目录已创建：{output_dir}")

    with open(output_path, "w", encoding="utf-8") as json_file:
        json.dump(data, json_file, ensure_ascii=False, indent=4)
    print(f"统计结果已保存到 {output_path}")


if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input-path', type=str, default='../Labels', help='BDD format label path')
    parser.add_argument('--output-path', type=str, default='./Json/7+7+4statistics.json', help='saved statistic path')
    opt = parser.parse_args()

    dict_count = attribute_statistics()
    output_path = opt.output_path

    dict_save2json(dict_count, output_path)
    pprint(dict_count)
